import Numeric as N
import string
import decimal

class PlateError( Exception ):
    pass

class LayoutError( PlateError ):
    pass

class PlateArray( object ):
    """
    Manage results of a single x by y well plate. 
    """

    def __init__(self, shape=(8,12), labelx=None, labely=None ):
        """
        @param shape: dimensions of the plate [8 letters by 12 numbers]
        @type  shape: (int, int)
        @param labelx: labels for rows [letters a,b,c,d...]
        @type  labelx: [ int or str ]
        @param labelx: labels for columns [numbers 1,2,3...]
        @type  labelx: [ int or str ]
        """
        self.time = None
        self.repeat = None
        self.plate = None
        self.temp_start = None
        self.temp_stop = None
        self.barcode = None

        self.__rows = None
        self.__cols = None

        self.protocols = {}  ## {'label' : 'value' }

        self.values = {}   ## {'label' : Numeric.array( rows x columns ) }

        self.setLayout( shape, labelx, labely )
        

    def setLayout( self, shape, labelx=None, labely=None ):
        """
        Change the plate layout or labeling.
        """
        labelx = labelx or string.letters[ : shape[0] ]
        labely = labely or range(1, shape[1] + 1 )

        if len( labelx ) != shape[0] or len( labely ) != shape[1]:
            raise LayoutError, 'Labels and shape are not aligned.'

        self.labelx = labelx
        self.labely = labely
        self.shape = shape
        

    def reIndex( self ):
        """
        Map x and y labels to array position.
        @return: two dictionaries -- row index & column index
        @rtype: ( { str : int }, { int : int } )
        """
        self.__rows, self.__cols = {}, {}
        
        for k, i in zip( self.labelx, range(len(self.labelx)) ) :
            self.__rows[ k ] = i

        for k, i in zip( self.labely, range(len(self.labely)) ) :
            self.__cols[ k ] = i


    def getRowIndex( self ):

        if self.__rows is None:
            self.reIndex()
        
        return self.__rows

    def getColIndex( self ):
        
        if self.__cols is None:
            self.reIndex()
        
        return self.__cols

    rowIndex = property( getRowIndex,
                     doc='map well row labels to array positions')

    colIndex = property( getColIndex,
                     doc='map well column labels to array positions')


    def values2array( self, values, wells ):
        """
        Convert list of values and associated well positions into a 2-D
        array.
        @param values: list or 1-D array of values
        @type  values: [ any ] or Numeric.array
        @param wells: list of well positions corresponding to each value
        @type  wells: [ (str,int) ], list of tuples like ('a',1)
        @return: array of values with self.shape 
        """
        assert len( values ) == len( wells ),\
               'A single well position is needed for each value.'

        s = self.shape

        r = N.zeros( s[0] * s[1], 'd' )

        ## translate well positions like ('a',1) to 1-D array position like 0
        rows, cols = self.rowIndex, self.colIndex

        valuepositions = [ (rows[w[0]] * s[1] + cols[w[1]] % s[0])
                           for w in wells ]
        
        N.put( r, valuepositions, values )

        return N.reshape( r, s )
    
            
    def get( self, name, default=None ):

        return self.values.get( name, default=default )


    def setWells( self, label, values, wells ):
        """
        Add or override a set of measurements.
        values - [ number ]; list or 1D array of measurements
        wells  - [ (str,int) ]; list of well positions (one for each value)
        """
        self.values[ label ] = self.values2array( values, wells )


    def getWells( self, wells ):
        """
        """
        pass
    

    def take( self, indices, axis=0 ):
        """
        @param indices: row or column positions
        @type  indices: [ int ]
        @param axis: take rows (0) or columns (1)
        @type  axis: int
        @return: new PlateArray with reduced or reordered layout
        @rtype: PlateArray
        """
        r = PlateArray()
        r.__dict__.update( self.__dict__ )
        r.values = {}

        shape = list( self.shape )
        shape[ axis ] = len( indices )
        shape = tuple( shape )

        labelx = self.labelx
        labely = self.labely
        if axis == 0:
            labelx = N.take( labelx, indices ).tolist()
        else:
            labely = N.take( labely, indices ).tolist()
            
        r.setLayout( shape, labelx, labely )

        for k in self.values.keys():
            r.values[ k ] = N.take( self.values[ k ], indices, axis=axis )

        return r


    def well2index( self, well ):
        """
        Translate well label into array position.
        """
        return self.rows[ well[0] ], self.cols[ well[1] ]

    def index2well( self, pos ):
        """
        Translate array position into well label.
        """
        pass


    def __getitem__( self, k ):
        """
        
        """
        if type( k ) is tuple:

            if len( k ) == 3:

                label, row, col = k
                x, y = self.rowIndex[ row ], self.colIndex[ col ]
                
                return self.values[ label ][x,y]

            if len( k ) == 2 and type( k[1] ) is str:

                label, row = k
                return N.take( self.values[ label ], self.rowIndex[ row ] )

            if len( k ) == 2 and type( k[1] ) is int:

                label, col = k
                return N.take( self.values[ label ], self.colIndex[ col ],
                               axis=1 )
        
        return self.values[ k ]

    
